import torch
import torchvision


def deeplabv3plus(classes):
    model = torchvision.models.segmentation.deeplabv3_resnet50(
                                                               num_classes=classes,
                                                               aux_loss=False)
    return model

import torch.nn as nn 

class DeepLabV3(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.model = deeplabv3plus(1)

    def forward(self, x):

        return self.model(x)["out"]



if __name__ == '__main__':
    model = deeplabv3plus(classes=1)
    print(model)
